package com.gitee.wsl.data.state.repo

import com.gitee.wsl.data.state.HandleAfterCreating
import com.gitee.wsl.data.state.HandleAfterNewStateIsSet
import com.gitee.wsl.data.state.HandleAfterOldStateIsRemoved
import com.gitee.wsl.data.state.HandleBeforeCreating
import com.gitee.wsl.data.state.HandleBeforeOldStateIsRemoved
import com.gitee.wsl.data.state.KalugaState
import com.gitee.wsl.ext.K
import com.gitee.wsl.ext.runBlocking
import kotlinx.atomicfu.atomic
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.CoroutineName
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.FlowCollector
import kotlinx.coroutines.flow.MutableSharedFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.first
import kotlinx.coroutines.launch
import kotlinx.coroutines.sync.Semaphore
import kotlinx.coroutines.sync.withPermit
import kotlin.coroutines.CoroutineContext
import kotlin.reflect.KClass

/**
 * The state repo can change holds the current [KalugaState] (which can be accessed as a flow), and can be used to change the current state
 *
 * @param State the type of [KalugaState] represented by this repo.
 * @param F the type of [MutableSharedFlow] the state will be collected in.
 * @param coroutineContext the [CoroutineContext] used to create a coroutine scope for this state machine. Make sure that if you pass a coroutine context that has sequential execution if you do not want simultaneous state changes. The default Main dispatcher meets these criteria.
 */
abstract class StateRepo<State : KalugaState, F : MutableSharedFlow<State>>(
    coroutineContext: CoroutineContext = Dispatchers.Main.immediate
) : CoroutineScope by CoroutineScope(coroutineContext + CoroutineName("State Repo")),
    Flow<State> {

    override suspend fun collect(collector: FlowCollector<State>) = mutableFlow.collect(collector)

    /**
     * Semaphore used as a mutex for changing state.
     *
     * By default the single permit is acquired, only when the initial state is set a release is done.
     */
    private val stateMutex = Semaphore(1, 1)

    protected abstract val mutableFlow: F

    /**
     * Number of subscribers to the state repo
     */
    val subscriptionCount
        get() = mutableFlow.subscriptionCount

    private val initialized = atomic(false)

    internal open suspend fun initialize(initialValue: State? = null): State
    = if (initialized.compareAndSet(expect = false, update = true)) {
        (initialValue ?: initialValue()).also { value ->
            mutableFlow.emit(value)
            stateMutex.release() // release the initial permit held
        }
    } else {
        state()
    }

    /**
     * Gets the initial value of the repo
     * @return the initial [State] of the repo
     */
    abstract suspend fun initialValue(): State

    internal suspend fun state(): State {
        // TODO:
        // if this state machine is backed by a SharedFlow instead of a pure StateFlow this will suspend indefinitely if no state is set
        // this only occurs (normally) if the initial state is not set.
        //
        // Perhaps an alternate way of throwing the exception (e.g. a boolean flag for the first emit) would be good,
        // however this does not guard manipulation of the SharedFlow before it is passed (e.g. MutableStateFlow's default value)
        // or directly on the flow (currently exposed to subclasses as a protected field)
        //
        // in the meanwhile, there might also be legitimate use cases for suspending until first state.
        // So for now at least, we accept this possible deadlock
        //
        // The replay cache instantiation is somewhat heavy but there is no method to get only the the last entry in the cache
        // if no or a small replayCache is used it is not so bad. Also is a StateFlow is used, this problem does not occur.
        @Suppress("UNCHECKED_CAST")
        return (mutableFlow as? StateFlow<State>)?.value ?: mutableFlow.replayCache.lastOrNull() ?: mutableFlow.first()
    }

    /**
     * Peek the current state of the state machine. The current state could change immediately after it is returned.
     *
     * Also no state could be set yet, in which case this method will block. Use [useState] for a suspending alternative.
     *
     * If any actions are taken based on the current state that affect the state machine you should not use this method.
     *
     * If your code relies on the state not changing use [useState].
     * If you want to change the state based on the current state use [takeAndChangeState]
     *
     * @return the current [State] of the [StateRepo]
     */
    fun peekState() = K.runBlocking {
        initialize()
        state()
    }

    /**
     * Makes the current [State] available in [action]. The state is guaranteed not to change during the execution of [action].
     * This operation ensures atomic state observations, so the state will not change while the [action] is being executed.
     *
     * This method uses a separate coroutineScope, meaning it will suspend until all child Jobs are completed, including those that asynchronously call this method itself (however a different state might be current at that point).
     *
     * @param Result the type of the result to be returned
     * @param action the function for which will [State] receive the state, guaranteed to be unchanged for the duration of the function.
     * @return the [Result] returned by [action] for the [State]
     */
    suspend fun <Result> useState(action: suspend (State) -> Result) = coroutineScope {
        initialize()
        stateMutex.withPermit {
            val result = CompletableDeferred<Result>()
            launch {
                try {
                    val actionResult = action(state())
                    result.complete(actionResult)
                } catch (e: Throwable) {
                    if (e is CancellationException) {
                        throw e
                    }
                    result.completeExceptionally(e)
                }
            }
            return@coroutineScope result.await()
        }
    }

    /**
     * Launches in a given [CoroutineContext] and calls [action] on the current [State]. The state is guaranteed not to change during the execution of [action].
     * This operation ensures atomic state observations, so the state will not change while the [action] is being executed.
     *
     * @param Result the type of the result to be returned
     * @param context The [CoroutineContext] in which to use the state.
     * @param action The action to execute on the current [State]
     * @see [useState]
     */
    fun <Result> launchUseState(context: CoroutineContext = coroutineContext, action: suspend (State) -> Result) = launch(context) {
        useState(action)
    }

    /**
     * Changes from the current [State] to a new [State]. This operation ensures atomic state changes.
     * The new state is determined by an [action], which takes the current [State] upon starting the state transition and provides a deferred state creation.
     * You are strongly encouraged to use the [State] provided by the [action] to determine the new state, to ensure no illegal state transitions occur, as the state may have changed between calling [doTakeAndChangeState] and the execution of [action].
     * If the [action] returns [KalugaState.remain] no state transition will occur.
     * Since this operation is atomic, the [action] should not directly call [doTakeAndChangeState] itself. If required to do this, handle the additional transition in a separate coroutine.
     *
     * This method uses a separate coroutineScope, meaning it will suspend until all child Jobs are completed, including those that asynchronously call this method itself.
     *
     * @param action Function to determine the [State] to be transitioned to from the current [State]. If no state transition should occur, return [KalugaState.remain]
     */
    suspend fun takeAndChangeState(action: suspend (State) -> suspend () -> State) = doTakeAndChangeState(remainIfStateNot = null, action)

    /**
     * Changes from the current [State] to a new [State]. This operation ensures atomic state changes.
     * The new state is determined by an [action], which takes the current [State] upon starting the state transition and provides a deferred state creation.
     * You are strongly encouraged to use the [State] provided by the [action] to determine the new state, to ensure no illegal state transitions occur, as the state may have changed between calling [doTakeAndChangeState] and the execution of [action].
     * If the [action] returns [KalugaState.remain] no state transition will occur.
     * Since this operation is atomic, the [action] should not directly call [doTakeAndChangeState] itself. If required to do this, handle the additional transition in a separate coroutine.
     *
     * This method uses a separate coroutineScope, meaning it will suspend until all child Jobs are completed, including those that asynchronously call this method itself.
     *
     * @param remainIfStateNot If the current state at the time of Action is not an instance of this class, the state will automatically remain.
     * @param action Function to determine the [State] to be transitioned to from the current [State]. If no state transition should occur, return [KalugaState.remain]
     */
    suspend fun <K : State> takeAndChangeState(remainIfStateNot: KClass<K>, action: suspend (K) -> suspend () -> State) = doTakeAndChangeState(remainIfStateNot) {
        action(it)
    }

    /**
     * Launches in a given [CoroutineContext] to change from the current [State] to a new [State]. This operation ensures atomic state changes.
     *
     * @param context The [CoroutineContext] in which to use the state.
     * @param action Function to determine the [State] to be transitioned to from the current [State]. If no state transition should occur, return [KalugaState.remain]
     * @see [takeAndChangeState]
     */
    fun launchTakeAndChangeState(context: CoroutineContext = coroutineContext, action: suspend (State) -> suspend () -> State) = launch(context) {
        takeAndChangeState(action)
    }

    /**
     * Launches in a given [CoroutineContext] to change from the current [State] to a new [State]. This operation ensures atomic state changes.
     *
     * @param context The [CoroutineContext] in which to use the state.
     * @param remainIfStateNot If the current state at the time of Action is not an instance of this class, the state will automatically remain.
     * @param action Function to determine the [State] to be transitioned to from the current [State]. If no state transition should occur, return [KalugaState.remain]
     * @see [takeAndChangeState]
     */
    fun <K : State> launchTakeAndChangeState(
        context: CoroutineContext = coroutineContext,
        remainIfStateNot: KClass<K>,
        action: suspend (K) -> suspend () -> State)
    = launch(context) {
            takeAndChangeState(remainIfStateNot) {
                action(it)
            }
        }

    @Suppress("UNCHECKED_CAST")
    private suspend inline fun <K : State> doTakeAndChangeState(
        remainIfStateNot: KClass<K>?,
        crossinline action: suspend (K) -> suspend () -> State): State =
        coroutineScope {
            // scope around the mutex so asynchronously scheduled coroutines that also use this method can run before the scope completed without deadlocks
            initialize()
            stateMutex.withPermit {
                val result = CompletableDeferred<State>()
                launch {
                    try {
                        val beforeState = state()
                        // There are only two methods calling this private method.
                        // either K is the same as S (no `remainIfStateNot` parameter), or we do the isInstance check
                        val transition =
                            // if remainIfNot was passes, only execute action if the beforeState matches
                            when {
                                remainIfStateNot == null || remainIfStateNot.isInstance(beforeState) ->
                                    action(beforeState as K)

                                else -> beforeState.remain() // else just remain
                            }

                        if (beforeState.remain<State>() === transition) {
                            result.complete(beforeState)
                        } else {
                            (beforeState as? HandleBeforeCreating)?.beforeCreatingNewState()
                            val newState = transition()
                            (beforeState as? HandleAfterCreating<State>)?.afterCreatingNewState(
                                newState
                            )
                            (newState as? HandleBeforeOldStateIsRemoved<State>)?.beforeOldStateIsRemoved(
                                beforeState
                            )
                            mutableFlow.emit(newState)
                            (beforeState as? HandleAfterNewStateIsSet<State>)?.afterNewStateIsSet(
                                newState
                            )
                            (newState as? HandleAfterOldStateIsRemoved<State>)?.afterOldStateIsRemoved(
                                beforeState
                            )
                            result.complete(newState)
                        }
                    } catch (t: Throwable) {
                        result.completeExceptionally(t)
                    }
                }
                return@coroutineScope result.await()
            }
        }
}

